'''利用RWKV模型处理用户输出的prompt，从而达到一定程度上的画面效果优化。

为了运行这个脚本，你需要设置环境变量：FAKE_PHOTOS_RWKV_MODEL_PATH，设置为你本地RWKV模型
的路径。比如，D:\RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth，那么设置：
FAKE_PHOTOS_RWKV_MODEL_PATH="D:\RWKV-x060-World-1B6-v2.1-20240328-ctx4096"

除了环境变量，命令行的参数也需要注意：

 -p 提示词
 -o 模型输出到此文件
 -s 随机数种子

示例：

export FAKE_PHOTOS_RWKV_MODEL_PATH="D:\RWKV-x060-World-1B6-v2.1-20240328-ctx4096"
$ python rwkv_worker.py -p "霞光照射在高楼上" -o b.txt -s 42
output file: b.txt
prompt: 霞光照射在高楼上
seed: 42
RWKV_JIT_ON 1 RWKV_CUDA_ON 0 RESCALE_LAYER 0

Loading D:\RWKV-x060-World-1B6-v2.1-20240328-ctx4096.pth ...
......
head.weight                       f32      cpu   2048 65536
Instruction: 按以下内容用英文描述画面

Input: 霞光照射在高楼上

Response: The sunlight illuminates the high buildings.

b.txt是模型的输出：
$ cat b.txt
The sunlight illuminates the high buildings.
'''

import random
import numpy as np
import torch
import re
import os
import sys
import argparse
from ChatRWKV import ChatRWKV

FAKE_PHOTOS_RWKV_MODEL_PATH = os.environ.get('FAKE_PHOTOS_RWKV_MODEL_PATH')
if FAKE_PHOTOS_RWKV_MODEL_PATH is None:
    print('FAKE_PHOTOS_RWKV_MODEL_PATH is not set!', file=sys.stderr)
    exit(1)

parser = argparse.ArgumentParser(description='Prompt processing service based on RWKV')
parser.add_argument('-o', help='output file')
parser.add_argument('-p', help='prompt')
parser.add_argument('-s', help='seed')

args, _ = parser.parse_known_args()
seed = random.randint(1, 2**16) if int(args.s) < 0 else int(args.s)
print(f"output file: {args.o}")
print(f"prompt: {args.p}")
print(f"seed: {seed}")

torch.manual_seed(seed)
np.random.seed(seed)

def model(inp):

    # 过滤所有换行符为空格
    prompt = re.sub(r'[\s|\n]+', ' ', inp)

    inputString = "Instruction: 按以下内容用英文描述画面\n\n"
    inputString = inputString + f"Input: {prompt}\n\n"
    inputString = inputString + "Response:"

    rwkv = ChatRWKV(FAKE_PHOTOS_RWKV_MODEL_PATH)

    print(inputString, end='')

    modelOutput = rwkv.run_on(
        inputString,
        temperature=1,
        top_p=0.3,
        top_k=15,
        max_length=800
    )
    rwkv.clear()
    return modelOutput

modelOutput = model(args.p)
print(' ' + modelOutput)

with open(args.o, 'w') as f:
    f.write(modelOutput)
